
import os
import torch
import torch.nn.functional as F
import torchtext.datasets as datasets

from torchtext.data.functional import to_map_style_dataset
from torch.utils.data import DataLoader
from torchtext.vocab import build_vocab_from_iterator
from torch.utils.data.distributed import DistributedSampler


VOCAB_PATH = "vocab.pt"



def tokenize(text, tokenizer):
    return [tok.text for tok in tokenizer.tokenizer(text)]


# def build_vocabulary_from_


# 根据训练数据来构建一张词汇表
# 接收两个参数：spacy_de 和 spacy_en，通常是用于德语和英语分词的 spaCy 分词器
def build_vocabulary(spacy_de, spacy_en):
    def tokenize_de(text):
        return tokenize(text, spacy_de)

    def tokenize_en(text):
        return tokenize(text, spacy_en)

    # 处理 德语 部分
    print("Building German Vocabulary ...")
    # 从 torchtext.datasets.Multi30k 加载德英平行语料数据（训练集、验证集、测试集）
    # 返回一个 Vocab 对象（torchtext.vocab.Vocab）
    # 可以 1.查词的 ID（vocab["word"]） 2.查 ID 对应的词（vocab.lookup_token(index)）
    # 3.设置默认索引（如 vocab.set_default_index(vocab["<unk>"])）
    train = datasets.Multi30k(split='train', language_pair=('de', 'en'))
    val   = datasets.Multi30k(split='valid', language_pair=('de', 'en'))
    test  = datasets.Multi30k(split='test', language_pair=('de', 'en'))

    # for example. next(iter(train)):
    # ('Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.',
    # 'Two young white males are outside near many bushes.')

    def yield_tokens(data_iter, tokenizer, index):
        for from_to_tuple in data_iter:
            yield tokenizer(from_to_tuple[index])

    # build_vocab_from_iterator 要求第一个参数也就是迭代器
    # 每次迭代返回的必须是一个 List[str]，也就是 token 列表
    vocab_src = build_vocab_from_iterator(
        # train + val + test 本质上有点类似于 元素是 tuple 的 list
        yield_tokens(train + val + test, tokenize_de, index=0),
        min_freq=2, # min_freq=2 表示只保留在语料中至少出现 2 次的词，低于 2 次的都 pass 掉
        specials=["<s>", "</s>", "<blank>", "<unk>"], # specials 是特殊符号列表：起始符、结束符、占位符、未知词
    )
    # specials 的作用就是向这个词汇表中添加一些字符串，这些字符串有固定的索引


    # 处理英语部分
    print("Building English Vocabulary ...")
    train = datasets.Multi30k(split='train', language_pair=('de', 'en'))
    val   = datasets.Multi30k(split='valid', language_pair=('de', 'en'))
    test  = datasets.Multi30k(split='test', language_pair=('de', 'en'))
    vocab_tgt = build_vocab_from_iterator(
        yield_tokens(train + val + test, tokenize_en, index=1),
        min_freq=2,
        specials=["<s>", "</s>", "<blank>", "<unk>"],
    )

    # 设置词汇表中未出现的词的默认索引为 <unk> 对应的索引
    vocab_src.set_default_index(vocab_src["<unk>"])
    vocab_tgt.set_default_index(vocab_tgt["<unk>"])

    return vocab_src, vocab_tgt


# ***************************************************************************************
# Usage : 加载 VOCAB_PATH 中的词汇表，如果没有这个词汇表的话，就自动调用 build_vocabulary 构建
#
# @input_param: 见函数开头
#
# @ret: 一个 tuple (vocab_src, vocab_tgt)，其中的 vocab_src, vocab_tgt 都是词汇表
#       所谓的 词汇表 就是 TODO:
# =======================================================================================
def load_vocab(spacy_de, spacy_en):
    if not os.path.exists(VOCAB_PATH):
        vocab_src, vocab_tgt = build_vocabulary(spacy_de, spacy_en)
        torch.save((vocab_src, vocab_tgt), VOCAB_PATH)
    else:
        vocab_src, vocab_tgt = torch.load(VOCAB_PATH)
    print("Finished.\nVocabulary sizes:")
    print(len(vocab_src))
    print(len(vocab_tgt))
    return (vocab_src, vocab_tgt)


# ***************************************************************************************
# Usage : 把一批 batch 的数据（元组列表）转成 一个 tuple（这个 tuple 中是 src 和 tgt 的组合）
#
# batch = [                     src_tensor = [             tgt_tensor = [
# ("I love you", "我 爱 你")       [0, 3, 4, 5, 1, 2],         [0, 10, 11, 12, 1, 2],
# ("Hello", "你好")                [0, 6, 1, 2, 2, 2],         [0, 13, 1, 2, 2, 2],
# ("Good morning", "早上 好")      [0, 7, 8, 1, 2, 2],         [0, 14, 15, 1, 2, 2],
# ]                             ]                          ]
#
# @input_param: 见函数开头
#
# @ret: 一个 tuple (src, tgt)，其中 src, tgt 都是 tensor, 形状是 (batch_size, max_padding)
#       (src_tensor(batch_size, max_padding), tgt_tensor(batch_size, max_padding))
# =======================================================================================
def collate_batch(
    batch,              # 一批数据，是一个 list,里面的元素是一个 tuple (src_sentence, tgt_sentence)
    src_pipeline,       # 源语言的分词函数（比如 str -> List[str]）
    tgt_pipeline,       # 目标语言的分词函数（比如 str -> List[str]）
    src_vocab,          # 源语言词表（把 token 列表转为 id 列表）
    tgt_vocab,          # 目标语言词表（把 token 列表转为 id 列表）
    device,             # 指定放到哪个设备（cpu/gpu）
    max_padding=128,    # 每个句子最多 pad 到多长
    pad_id=2,           # padding 的 token id（对应 <blank>）
):
    bs_id = torch.tensor([0], device=device)  # <s> token id
    eos_id = torch.tensor([1], device=device)  # </s> token id

    # bs_id = torch.tensor([src_vocab["<s>"]], device=device)
    # eos_id = torch.tensor([src_vocab["</s>"]], device=device)

    # 初始化两个空列表，src_list 用来存储处理后的源语言句子张量，
    # tgt_list 用来存储处理后的目标语言句子张量。
    src_list, tgt_list = [], []

    # batch 是一个列表，里面每个元素是 (source_sentence_str, target_sentence_str) 的元组。
    for (_src, _tgt) in batch:
        # _src 被 src_pipeline 处理成一个 token 列表。
        # 然后 src_vocab(...) 把 token 列表 转成 ID 列表
        # 接着 用 torch.tensor(...) 将 ID 列表 变成张量
        # 用 torch.cat(...) 添加起始符 <s> 和终止符 </s>。
        processed_src = torch.cat(
            [
                bs_id,
                torch.tensor(
                    src_vocab(src_pipeline(_src)),
                    dtype=torch.int64,
                    device=device,
                ),
                eos_id,
            ],
            0,
        )
        processed_tgt = torch.cat(
            [
                bs_id,
                torch.tensor(
                    tgt_vocab(tgt_pipeline(_tgt)),
                    dtype=torch.int64,
                    device=device,
                ),
                eos_id,
            ],
            0,
        )
        # 句子太长的话，要进行截断
        if len(processed_src) > max_padding:
            print(
                f"⚠️ Warning: source sequence is too long "
                f"(length={len(processed_src)}), truncating to {max_padding}."
            )
            processed_src = processed_src[:max_padding]

        src_list.append(
            # 把 src padding 到 max_padding
            F.pad(processed_src, (0, max_padding - len(processed_src), ), value=pad_id)
        )
        # 句子太长的话，要进行截断
        if len(processed_tgt) > max_padding:
            print(
                f"⚠️ Warning: target sequence is too long "
                f"(length={len(processed_tgt)}), truncating to {max_padding}."
            )
            processed_tgt = processed_tgt[:max_padding]

        tgt_list.append(
            # 把 tgt padding 到 max_padding
            F.pad(processed_tgt, (0, max_padding - len(processed_tgt)), value=pad_id)
        )
    # end of 'for'

    src = torch.stack(src_list)
    tgt = torch.stack(tgt_list)
    return (src, tgt)


# ***************************************************************************************
# Usage : 创建 train 和 valid 的 dataloaders
#
# @input_param: 见函数开头
#
# @ret: 一个 tuple (train_dataloader, valid_dataloader)
# =======================================================================================
def create_dataloaders(
    device,               # 在 cpu 还是在 gpu 上面进行训练
    vocab_src,            # 源语言的词表，用于把单词转成对应的 id
    vocab_tgt,            # 目标语言的词表，用于把单词转成对应的 id
    spacy_de,             # 德语的 spaCy 分词模型对象
    spacy_en,             # 英语的 spaCy 分词模型对象
    batch_size=12000,     # 每个批次多少条样本，默认12000（比较大，视显存和任务调整）
    max_padding=128,      # padding最大长度，句子会pad到这个长度
    is_distributed=True,  # 是否使用分布式训练，决定是否用分布式采样器
):
    def tokenize_de(text):
        return tokenize(text, spacy_de)

    def tokenize_en(text):
        return tokenize(text, spacy_en)

    def collate_fn(batch):
        return collate_batch(
            batch,
            tokenize_de,
            tokenize_en,
            vocab_src,
            vocab_tgt,
            device,
            max_padding=max_padding,
            pad_id=vocab_src.get_stoi()["<blank>"],
        )

    train_iter = datasets.Multi30k(split='train', language_pair=('de', 'en'))
    valid_iter = datasets.Multi30k(split='valid', language_pair=('de', 'en'))

    train_iter_map = to_map_style_dataset(
        train_iter
    )  # DistributedSampler needs a dataset len()
    train_sampler = (
        DistributedSampler(train_iter_map) if is_distributed else None
    )
    valid_iter_map = to_map_style_dataset(valid_iter)
    valid_sampler = (
        DistributedSampler(valid_iter_map) if is_distributed else None
    )
    # 创建 训练集 的 dataloader
    train_dataloader = DataLoader(
        train_iter_map,
        batch_size=batch_size,
        shuffle=(train_sampler is None),  # 打乱最底层的数据集中的数据顺序
        sampler=train_sampler,
        collate_fn=collate_fn,  # 自定义如何把多个样本组成一个 batch。默认是把 list 转成 tensor
    )
    # 创建 验证集 的 dataloader
    valid_dataloader = DataLoader(
        valid_iter_map,
        batch_size=batch_size,
        shuffle=(valid_sampler is None),  # 打乱最底层的数据集中的数据顺序
        sampler=valid_sampler,
        collate_fn=collate_fn,  # 自定义如何把多个样本组成一个 batch。默认是把 list 转成 tensor
    )
    return train_dataloader, valid_dataloader


if __name__ == "__main__":
    train_dataloader, valid_dataloader = create_dataloaders()

